#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件        :    make_coco_mini.py
@说明        :    
@时间        :    2021/12/21 14:33:43
@作者        :    钟建平
@版本        :    1.0
'''

import json
import argparse
import os
import numpy as np
from tqdm import tqdm
import shutil

def  _parser_args():
    """parse input"""
    parser = argparse.ArgumentParser("convert coco to coco mini")
    parser.add_argument('--coco_path', type=str, required=True, help='the path of coco')
    parser.add_argument('--ratio', type=float, required=True, help='the path of coco')
    parser.add_argument('--coco_mini_path', type=str,  required=True,  help='the path ot save coco_mini data')
    args, _ = parser.parse_known_args()
    return args

def split(data_type, coco_path, coco_mini_path, ratio):
    print(f'--------------start to process {data_type}--------------------')
    image_id = []
    result = dict()
    src = json.load(open(coco_path+f'annotations/instances_{data_type}.json'))
    result['annotations'] = []
    result['images'] = []
    result['categories'] = src['categories']
    src_image = src['images']
    np.random.shuffle(src_image)
    for i in tqdm(range(int(ratio*len(src_image)))):
        result['images'].append(src_image[i])
        image_id.append(src_image[i]['id'])
        image_name = src_image[i]['file_name']
        shutil.copy(coco_path+data_type+'/'+image_name, coco_mini_path+data_type+'/'+image_name)

    for annotation in tqdm(src['annotations']):
        if annotation['image_id'] in image_id:
            result['annotations'].append(annotation)
    with open(coco_mini_path+f'annotations/instances_{data_type}.json', 'w') as f:
        json.dump(result, f, indent='\t')
    print(f'--------------process {data_type} finished!--------------------')
def make(coco_path, ratio, coco_mini_path):
    if not os.path.exists(coco_mini_path):
        os.mkdir(coco_mini_path)
    if not os.path.exists(coco_mini_path+'annotations'):
        os.mkdir(coco_mini_path+'annotations')
    if not os.path.exists(coco_mini_path+'train2017'):
        os.mkdir(coco_mini_path+'train2017')
    if not os.path.exists(coco_mini_path+'val2017'):
        os.mkdir(coco_mini_path+'val2017')
    split('train2017', coco_path,  coco_mini_path, ratio)
    split('val2017', coco_path, coco_mini_path, ratio)
    

if __name__ == '__main__':
    args = _parser_args()
    make(args.coco_path, args.ratio, args.coco_mini_path)
